# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

import logging
import re

from experimental.prophet.kernels import op_kernel
from experimental.prophet.op_registry import register_op

from coin.base.datetime_util import to_datetime


def _identity(x):
  return x


class Printer(op_kernel.OpKernel):
  allow_null_input = True

  name = 'Printer'
  input_args = op_kernel.variadic_input()
  output_args = []

  def __init__(self, attrs, node_def):
    format_str = attrs['format_str']
    self._fmt = []
    for m in re.finditer('%.*?([stbcdixXneEfFgGn%])', format_str):
      if m.group(1) == '%':
        continue
      if m.group(1) == 't':
        eidx = m.span()[1]
        format_str = format_str[:eidx - 1] + 's' + format_str[eidx:]
        self._fmt.append(to_datetime)
      else:
        self._fmt.append(_identity)

    while len(node_def.inputs) > len(self._fmt):
      self._fmt.append(_identity)

    self._format_str = format_str
    self._use_logger = attrs.get('use_logger', False)

  def compute(self, inputs, outputs):
    values = tuple([
        self._fmt[idx](var.value) if var.value is not None else None for idx,
        var in enumerate(inputs)
    ])

    if self._use_logger:
      logging.info(self._format_str, *values)
    else:
      print(self._format_str % values)


register_op(Printer)
